from paddleocr import PaddleOCR, draw_ocr
import fitz
import math
from datetime import datetime

start_time = datetime.now()


# Paddleocr目前支持的多语言语种可以通过修改lang参数进行切换
# 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`
# need to run only once to download and load model into memory
ocr = PaddleOCR(use_angle_cls=True, lang="ch", page_num=0, use_gpu=False)
img_path = r'/home/www/xerdoc70/public/tools/ppocr/test/input.pdf'
output_path = r"/home/www/xerdoc70/public/tools/ppocr/test/output.pdf"
font_path = r'D:\work\PythonProject\XerDocPaddleOcr\ppocr_img\fonts/simfang.ttf'
result = ocr.ocr(img_path, cls=True)
for idx in range(len(result)):
    res = result[idx]
    for line in res:
        print(line)


# 将识别的文本作为透明文本覆盖到原来的pdf文件中
with fitz.open(img_path) as pdf:
    for pg in range(0, pdf.page_count):
        page = pdf[pg]
        res = result[pg]

        # 因为ppocr采用pdf页面长宽小于1000而采用放大策略，所以需要考虑缩放尺寸
        scale = 0.5
        pm = page.get_pixmap(alpha=False)
        # if width or height > 2000 pixels, don't enlarge the image
        if pm.width > 1000 or pm.height > 1000:
            scale = 1.0

        # 将识别的文本框画到pdf页面上
        for box, [txt, score] in res:
            print(box, txt, score)
            # 文字框的四边形坐标，因为ocr翻了倍数2，所以需要缩小到0.5倍
            quad = fitz.Quad(box[0], box[1], box[2], box[3]) * scale
            # 计算四边形近似的宽高
            width = quad.ul.distance_to(quad.ur)
            height = quad.ul.distance_to(quad.lr)

            # 得到水平文字框矩形以及旋转角度
            rect = None
            d = None
            if len(txt) > 1 > width/height:
                rect = fitz.Rect(quad.ur, quad.ur.x + height, quad.ur.y + width)
                d = quad.ll - quad.ur
            else:
                rect = fitz.Rect(quad.ul, quad.ul.x + width, quad.ul.y + height)
                d = quad.ur - quad.ul

            # 使用math.atan2计算夹角（弧度）
            angle_rad = math.atan2(d.y, d.x)
            # 将弧度转换为角度
            angle_deg = math.degrees(angle_rad)

            # 计算文本框的旋转中心点位，及角度
            morph = (rect.tl, fitz.Matrix(1, 1).prerotate(-angle_deg))
            fontsize = 50
            rc = -1
            while rc < 0:  # look for largest font size that makes the text fit
                rc = page.insert_textbox(rect, txt,
                                         fontname="china-ss",
                                         fontsize=fontsize,
                                         color=(0, 0, 0, 0),
                                         fill_opacity=0.01,
                                         lineheight=quad.height,
                                         morph=morph,
                                         )
                fontsize -= 0.2  # 步进减小字体大小，直到匹配框大小
        # page.commit()
    pdf.save(output_path)

end_time = datetime.now()
time_diff = end_time - start_time
print(time_diff)